import json
import random
from synonyms import synonyms
import numpy as np

# 同义词替换
def synonym_replacement(text, n=1):
    words = text.split()
    if len(words) < 2:
        return text
        
    new_words = words.copy()
    random_word_list = list(set([word for word in words if word not in ['，', '。', '的', '是']]))
    random.shuffle(random_word_list)
    
    num_replaced = 0
    for random_word in random_word_list:
        try:
            syns = synonyms.nearby(random_word)[0]
            if len(syns) >= 1:
                synonym = random.choice(syns)
                new_words = [synonym if word == random_word else word for word in new_words]
                num_replaced += 1
            if num_replaced >= n:
                break
        except:
            continue
            
    return ' '.join(new_words)

# 随机插入
def random_insertion(text, n=1):
    words = text.split()
    if len(words) < 2:
        return text
        
    new_words = words.copy()
    for _ in range(n):
        random_word = random.choice(words)
        try:
            syns = synonyms.nearby(random_word)[0]
            if len(syns) > 0:
                synonym = random.choice(syns)
                insert_pos = random.randint(0, len(new_words))
                new_words.insert(insert_pos, synonym)
        except:
            continue
            
    return ' '.join(new_words)

# 数据增强主函数
def main():
    # 加载处理后的数据
    with open('../data/processed/train.json', 'r', encoding='utf-8') as f:
        train_data = json.load(f)
    
    # 创建增强数据集
    augmented_data = []
    for item in train_data:
        # 对每条数据生成3种增强版本
        augmented_texts = [
            synonym_replacement(item['text']),
            random_insertion(item['text']),
            synonym_replacement(random_insertion(item['text']))
        ]
        
        for text in augmented_texts:
            augmented_data.append({
                'text': text,
                'label': item['label'],
                'intent': item['intent']
            })
    
    # 保存增强数据
    with open('../data/augmented/augmented_train.json', 'w', encoding='utf-8') as f:
        json.dump(augmented_data, f, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    main()